# -*- coding: utf-8 -*-
"""
中性策略框架 | 邢不行 | 2024分享会
author: 邢不行
微信: xbx6660
"""
import os
import pandas as pd

# 取出策略文件名中_后的部分
stg_name = os.path.basename(__file__).split('.')[0].split('_')[1]

_ = os.path.abspath(__file__)
stg_folder = os.path.basename(os.path.dirname(_))
base_stg_name = os.path.basename(_).replace(".py", "")

# 分时间段定义的子策略字典
# bys_rolling_stg_dict = {
#     "子策略_001": {"起始日期": None, "结束日期": "2017-12-31"},
#     "子策略_002": {"起始日期": "2018-01-01", "结束日期": "2019-12-31"},
#     "子策略_003": {"起始日期": "2020-01-01", "结束日期": None},
# }

# @$开始定义滚动炼丹结果数据$@#
bys_rolling_stg_dict = {}
# @$结束定义滚动炼丹结果数据$@#

if len(bys_rolling_stg_dict) == 0:
    raise Exception("未配置滚动子策略")

rolling_cls = {}
default_cls = None
for key in bys_rolling_stg_dict.keys():
    cls = __import__(
        f'program.{stg_folder}.{base_stg_name}_子策略.{key}', fromlist=("",)
    )
    rolling_cls[key] = cls
    if default_cls is None:
        default_cls = cls

# 持仓周期。目前回测支持日线级别、小时级别。例：1H，6H，3D，7D......
# 当持仓周期为D时，选币指标也是按照每天一根K线进行计算。
# 当持仓周期为H时，选币指标也是按照每小时一根K线进行计算。
hold_period = default_cls.hold_period
# 配置offset
offset = default_cls.offset
# 是否使用现货
if_use_spot = default_cls.if_use_spot  # True：使用现货。False：不使用现货，只使用合约。

# 多头选币数量。1 表示做多一个币; 0.1 表示做多10%的币
long_select_coin_num = default_cls.long_select_coin_num
# 空头选币数量。1 表示做空一个币; 0.1 表示做空10%的币
# short_select_coin_num = 0.1
short_select_coin_num = default_cls.short_select_coin_num  # long_nums意为着空头数量和多头数量保持一致。最多为所有合约的数量。注意：多头为0的时候，不能配置'long_nums'

# 多头的选币因子列名。
long_factor = default_cls.long_factor  # 因子：表示使用复合因子，默认是 factor_list 里面的因子组合。需要修改 calc_factor 函数配合使用
# 空头的选币因子列名。多头和空头可以使用不同的选币因子
short_factor = default_cls.long_factor

# 选币因子信息列表，用于`2_选币_单offset.py`，`3_计算多offset资金曲线.py`共用计算资金曲线
factor_list = []
for key, cls in rolling_cls.items():
    factor_list.extend(cls.factor_list)
# 将权重统一改为1以去重
factor_list = [(factor[0], factor[1], factor[2], 1) for factor in factor_list]
factor_list = list(set(factor_list))

# 确认过滤因子及其参数，用于`2_选币_单offset.py`进行过滤
filter_list = []
for key, cls in rolling_cls.items():
    filter_list.extend(cls.filter_list)
filter_list = list(set(filter_list))


def after_merge_index(df, symbol, factor_dict, data_dict):
    """
    合并指数数据之后的处理流程，非必要。
    本函数住要的作用如下：
            1、指定K线上不常用的数据在resample时的规则，例如：'taker_buy_quote_asset_volume': 'sum'
            2、合并外部数据，并指定外部数据在resample时的规则。例如：（伪代码）
                    chain_df = pd.read_csv(chain_path + symbol)  # 读取指定币种链上数据
                    df = pd.merge(df,chain_df,'left','candle_begin_time') # 将链上数据合并到日线上
                    factor_dict['gas_total'] = 'sum' # 链上的gas_total字段，在小时数据转日线数据时用sum处理
                    data_dict['gas_total'] = 'sum' # 链上的gas_total字段，在最后一次resample中的处理规则
    :param df:
    :param symbol:
    :param factor_dict: 小时级别resample到日线级别时使用(计算日线级别因子时需要，计算小时级别因子时不需要)
    :param data_dict: resample生成最终数据时使用
    :return:
    """

    return default_cls.after_merge_index(df, symbol, factor_dict, data_dict)

def after_resample(df, symbol):
    """
    数据重采样之后的处理流程，非必要
    :param df:          传入的数据
    :param symbol:      币种名称
    :return:
    """

    return default_cls.after_resample(df, symbol)


# =====================以上是数据整理部分封装转的策略代码==========================
# ============================================================================
# ============================================================================
# ============================================================================
# ============================================================================
# =======================以下是选币函数封装的策略代码=============================


def calc_factor(df, **kwargs):
    """
    计算因子

    多空使用相同的两个因子

        1。多空相同的单因子选币，不需要编写代码

        2。多空相同的复合因子，本案例代码不需要修改，直接使用即可。《本案例代码就是多空相同的复合因子选币》

    多空使用不同的两个因子

        1。多空不同的单因子选币，不需要编写代码

        2。多空分离选币，需要根据具体因子具体改写
            !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            !!!!!这里改写需要一定的代码能力!!!!!
            !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            多空存在某一方是复合因子，另一方是单因子，则只需要计算一个复合因子即可。代码可以参考下方案例代码

            如果多空都是复合因子，多空需要分别是计算复合因子。

    :param df:          原数据
    :return:
    """
    df_list = []
    for n, (key, dt) in enumerate(bys_rolling_stg_dict.items()):
        cls = rolling_cls[key]
        con = True
        if n > 0 and dt["起始日期"] is not None:
            # 第一个子策略不过滤起始日期
            con &= (df["candle_begin_time"] >= pd.to_datetime(dt["起始日期"]))
        if n < len(bys_rolling_stg_dict) - 1 and dt["结束日期"] is not None:
            # 最后一个子策略不过滤结束日期
            con &= (df["candle_begin_time"] <= pd.to_datetime(dt["结束日期"]))
        if not isinstance(con, bool):
            a_df = df[con]
        else:
            a_df = df
        a_df = cls.calc_factor(a_df.copy(), **kwargs)
        df_list.append(a_df)

    all_data = pd.concat(df_list)

    return all_data


def before_filter(df, **kwargs):
    """
    前置过滤函数
    自定义过滤规则，可以对多空分别自定义过滤规则

    :param df:                  原始数据
    :return:                    过滤后的数据
    """
    df_long_list = []
    df_short_list = []
    for n, (key, dt) in enumerate(bys_rolling_stg_dict.items()):
        cls = rolling_cls[key]
        con = True
        if n > 0 and dt["起始日期"] is not None:
            # 第一个子策略不过滤起始日期
            con &= (df["candle_begin_time"] >= pd.to_datetime(dt["起始日期"]))
        if n < len(bys_rolling_stg_dict) - 1 and dt["结束日期"] is not None:
            # 最后一个子策略不过滤结束日期
            con &= (df["candle_begin_time"] <= pd.to_datetime(dt["结束日期"]))
        if not isinstance(con, bool):
            a_df = df[con]
        else:
            a_df = df
        df_long, df_short = cls.before_filter(a_df.copy(), **kwargs)
        df_long_list.append(df_long)
        df_short_list.append(df_short)

    df_long = pd.concat(df_long_list)
    df_short = pd.concat(df_short_list)

    return df_long, df_short
